package com.yoyosys.db.partition.query.impl;

import com.yoyosys.db.partition.constant.PartitionType;
import com.yoyosys.db.partition.query.PartitionQuery;
import com.yoyosys.db.partition.vo.Partition;
import com.yoyosys.db.partition.vo.PartitionInfo;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;

/**
 * MySQL 分区分桶信息查询器
 */
public class MySqlPartitionQuery implements PartitionQuery {
    private static MySqlPartitionQuery instance = new MySqlPartitionQuery();
    private MySqlPartitionQuery(){}
    public static MySqlPartitionQuery getInstance(){
        return instance;
    }

    /**
     * 查询表分区信息
     *
     * @param conn
     * @param schemaName 库名
     * @param tableName
     * @return
     */
    @Override
    public PartitionInfo getPartition(Connection conn, String schemaName, String tableName) throws SQLException {
        PartitionInfo info = new PartitionInfo();
        info.setSchemaName(schemaName);
        info.setTableName(tableName);
        info.setPartitions(new ArrayList<>());

        // 获取分区类型和键
        String partitionTypeSql = new StringBuilder()
                .append("SELECT PARTITION_METHOD, PARTITION_EXPRESSION ")
                .append("FROM INFORMATION_SCHEMA.PARTITIONS ")
                .append("WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? ")
                .append("LIMIT 1")
                .toString();
        System.out.println("partitionTypeSql: " + partitionTypeSql);
        try (PreparedStatement pstmt = conn.prepareStatement(partitionTypeSql)) {
            pstmt.setString(1, schemaName);
            pstmt.setString(2, tableName);
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    info.setPartitionType(rs.getString("PARTITION_METHOD"));
                    info.setPartitionKey(rs.getString("PARTITION_EXPRESSION"));
                } else {
                    info.setPartitionType(PartitionType.NONE.name());
                    return info;
                }
            }
        }

        // 获取所有分区
        String partitionsSql =  new StringBuilder()
                .append("SELECT PARTITION_NAME, PARTITION_EXPRESSION, ")
                .append("TABLE_ROWS, PARTITION_DESCRIPTION ")
                .append("FROM INFORMATION_SCHEMA.PARTITIONS ")
                .append("WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? ")
                .append("ORDER BY PARTITION_ORDINAL_POSITION")
                .toString();
        System.out.println("partitionsSql: " + partitionsSql);
        try (PreparedStatement pstmt = conn.prepareStatement(partitionsSql)) {
            pstmt.setString(1, schemaName);
            pstmt.setString(2, tableName);
            try (ResultSet rs = pstmt.executeQuery()) {
                while (rs.next()) {
                    Partition partition = new Partition();
                    partition.setName(rs.getString("PARTITION_NAME"));
                    partition.setExpression(rs.getString("PARTITION_EXPRESSION"));
                    partition.setRowCount(rs.getLong("TABLE_ROWS"));
                    partition.setDescription(rs.getString("PARTITION_DESCRIPTION"));
                    info.getPartitions().add(partition);
                }
            }
        }

        return info;
    }

}
